{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "91265a57",
   "metadata": {},
   "outputs": [],
   "source": [
    "from huggingface_hub import hf_hub_download\n",
    "\n",
    "f = hf_hub_download(\n",
    "    repo_id=\"mesolitica/MusicBench-Instructions\",\n",
    "    repo_type='dataset',\n",
    "    filename=\"data/train-00000-of-00001.parquet\",\n",
    "    local_dir=\"./MusicBench-Instructions\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "2c3ea8fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "from glob import glob\n",
    "from tqdm import tqdm\n",
    "from multiprocess import Pool\n",
    "import itertools\n",
    "import zipfile\n",
    "import os\n",
    "\n",
    "def chunks(l, n):\n",
    "    for i in range(0, len(l), n):\n",
    "        yield (l[i: i + n], i // n)\n",
    "\n",
    "\n",
    "def multiprocessing(strings, function, cores=6, returned=True):\n",
    "    df_split = chunks(strings, len(strings) // cores)\n",
    "    pool = Pool(cores)\n",
    "    pooled = pool.map(function, df_split)\n",
    "    pool.close()\n",
    "    pool.join()\n",
    "\n",
    "    if returned:\n",
    "        return list(itertools.chain(*pooled))\n",
    "\n",
    "def loop(files):\n",
    "    files, _ = files\n",
    "    for zip_file_path in tqdm(files):\n",
    "        destination_folder = './'\n",
    "        with zipfile.ZipFile(zip_file_path, 'r') as zip_ref:\n",
    "            zip_ref.extractall(destination_folder)\n",
    "        os.remove(zip_file_path)\n",
    "\n",
    "# files = glob('*.zip')\n",
    "# if len(files):\n",
    "#     multiprocessing(files, loop, cores = min(len(files), 20), returned = False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "50e40c88",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/lib/python3/dist-packages/scipy/__init__.py:146: UserWarning: A NumPy version >=1.17.3 and <1.25.0 is required for this version of SciPy (detected version 1.26.4\n",
      "  warnings.warn(f\"A NumPy version >={np_minversion} and <{np_maxversion}\"\n"
     ]
    }
   ],
   "source": [
    "from transformers import AutoProcessor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "afef4611",
   "metadata": {},
   "outputs": [],
   "source": [
    "processor = AutoProcessor.from_pretrained(\"Qwen/Qwen2-Audio-7B-Instruct\")\n",
    "tokenizer = processor.tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "1e4fea4f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'question': 'can u describe the audio',\n",
       " 'answer': 'This mellow instrumental track showcases a dominant electric guitar that opens with a descending riff, followed by arpeggiated chords, hammer-ons, and a slide. The percussion section keeps it simple with rim shots and a common time count, while the bass adds a single note on the first beat of every bar. Minimalist piano chords round out the song while leaving space for the guitar to shine. There are no vocals, making it perfect for a coffee shop or some chill background music. The key is in E major, with a chord progression that centers around that key and a straightforward 4/4 time signature.',\n",
       " 'audio_filename': 'datashare/data_aug2/-0SdAVK79lg_1.wav',\n",
       " 'metadata': '{\"dataset\": \"MusicBench\", \"location\": \"data_aug2/-0SdAVK79lg_1.wav\", \"main_caption\": \"This mellow instrumental track showcases a dominant electric guitar that opens with a descending riff, followed by arpeggiated chords, hammer-ons, and a slide. The percussion section keeps it simple with rim shots and a common time count, while the bass adds a single note on the first beat of every bar. Minimalist piano chords round out the song while leaving space for the guitar to shine. There are no vocals, making it perfect for a coffee shop or some chill background music. The key is in E major, with a chord progression that centers around that key and a straightforward 4/4 time signature.\", \"alt_caption\": \"This song features an electric guitar as the main instrument. The guitar plays a descending run in the beginning then plays an arpeggiated chord followed by a double stop hammer on to a higher note and a descending slide followed by a descending chord run. The percussion plays a simple beat using rim shots. The percussion plays in common time. The bass plays only one note on the first count of each bar. The piano plays backing chords. There are no voices in this song. The mood of this song is relaxing. This song can be played in a coffee shop. The key of this song is E major. The chord progression in this song is E. The beat counts to 4. \", \"prompt_aug\": \"\", \"prompt_ch\": \"The chord progression in this song is E.\", \"prompt_bt\": \"The beat counts to 4.\", \"prompt_bpm\": \"The bpm is 112.0.\", \"prompt_key\": \"The key of this song is E major.\", \"beats\": [[0.37212133669135816, 0.9070457581851855, 1.4652277632222228, 2.0234097682592598, 2.5583341897530874, 3.093258611246915, 3.628183032740742, 4.13984987069136, 4.6747742921851865, 5.232956297222224, 5.744623135172842, 6.326062723753089, 6.860987145246916, 7.372653983197533, 7.930835988234571, 8.442502826185187, 8.977427247679016, 9.535609252716052], [4.0, 1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0, 1.0]], \"bpm\": 112.0, \"chords\": [\"E\"], \"chords_time\": [0.5400400339790407], \"key\": [\"E\", \"major\"], \"keyprob\": [0.8934084177017212], \"is_audioset_eval_mcaps\": false}'}"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import pandas as pd\n",
    "\n",
    "df = pd.read_parquet(f).to_dict(orient = 'records')\n",
    "df[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "848bb96f",
   "metadata": {},
   "outputs": [],
   "source": [
    "conversation = [\n",
    "    {\"role\": \"user\", \"content\": [\n",
    "        {\"type\": \"audio\", \"audio_url\": \"audio.wav\"},\n",
    "        {\"type\": \"text\", \"text\": df[0]['question']},\n",
    "    ]},\n",
    "    {\"role\": \"assistant\", \"content\": df[0]['answer']},\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "9f6a860d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'<|im_start|>system\\nYou are a helpful assistant.<|im_end|>\\n<|im_start|>user\\nAudio 1: <|audio_bos|><|AUDIO|><|audio_eos|>\\ncan u describe the audio<|im_end|>\\n<|im_start|>assistant\\nThis mellow instrumental track showcases a dominant electric guitar that opens with a descending riff, followed by arpeggiated chords, hammer-ons, and a slide. The percussion section keeps it simple with rim shots and a common time count, while the bass adds a single note on the first beat of every bar. Minimalist piano chords round out the song while leaving space for the guitar to shine. There are no vocals, making it perfect for a coffee shop or some chill background music. The key is in E major, with a chord progression that centers around that key and a straightforward 4/4 time signature.<|im_end|>\\n'"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "processor.apply_chat_template(conversation, tokenize=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "9bb8c8b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import pandas as pd\n",
    "\n",
    "def loop(rows):\n",
    "    rows, _ = rows\n",
    "    data = []\n",
    "    for r in tqdm(rows):\n",
    "        f = r['audio_filename']\n",
    "        if not os.path.exists(f):\n",
    "            continue\n",
    "            \n",
    "        try:\n",
    "            conversation = [\n",
    "                {\"role\": \"user\", \"content\": [\n",
    "                    {\"type\": \"audio\", \"audio_url\": \"audio.wav\"},\n",
    "                    {\"type\": \"text\", \"text\": r['question']},\n",
    "                ]},\n",
    "                {\"role\": \"assistant\", \"content\": r['answer']},\n",
    "            ]\n",
    "            text = processor.apply_chat_template(conversation, tokenize=False)\n",
    "        except Exception as e:\n",
    "            continue\n",
    "        \n",
    "\n",
    "        data.append({\n",
    "            'text': text,\n",
    "            'audio': f,\n",
    "        })\n",
    "    return data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "aea82c85",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████████████████████████████████████████████| 5039/5039 [00:00<00:00, 5618.21it/s]\n",
      "100%|█████████████████████████████████████████████████████████████████████████████████| 5039/5039 [00:00<00:00, 5391.92it/s]\n",
      "100%|█████████████████████████████████████████████████████████████████████████████████| 5039/5039 [00:00<00:00, 5659.15it/s]\n",
      "100%|█████████████████████████████████████████████████████████████████████████████████| 5039/5039 [00:00<00:00, 5645.27it/s]\n",
      "100%|█████████████████████████████████████████████████████████████████████████████████| 5039/5039 [00:00<00:00, 5634.72it/s]\n",
      "100%|█████████████████████████████████████████████████████████████████████████████████| 5039/5039 [00:00<00:00, 5629.11it/s]\n",
      "100%|█████████████████████████████████████████████████████████████████████████████████| 5039/5039 [00:00<00:00, 5618.58it/s]\n",
      "100%|█████████████████████████████████████████████████████████████████████████████████| 5039/5039 [00:00<00:00, 5561.78it/s]\n",
      "100%|█████████████████████████████████████████████████████████████████████████████████| 5039/5039 [00:00<00:00, 5605.86it/s]\n",
      "100%|█████████████████████████████████████████████████████████████████████████████████| 5039/5039 [00:00<00:00, 5631.06it/s]\n",
      "100%|█████████████████████████████████████████████████████████████████████████████████| 5039/5039 [00:00<00:00, 5669.80it/s]\n",
      "100%|█████████████████████████████████████████████████████████████████████████████████| 5039/5039 [00:00<00:00, 5610.32it/s]\n",
      "100%|█████████████████████████████████████████████████████████████████████████████████| 5039/5039 [00:00<00:00, 5258.61it/s]\n",
      " 89%|████████████████████████████████████████████████████████████████████████▎        | 4496/5039 [00:00<00:00, 5724.99it/s]\n",
      "100%|█████████████████████████████████████████████████████████████████████████████████| 5039/5039 [00:00<00:00, 5625.99it/s]\n",
      "  0%|                                                                                              | 0/5039 [00:00<?, ?it/s]\n",
      "100%|█████████████████████████████████████████████████████████████████████████████████| 5039/5039 [00:00<00:00, 5614.13it/s]\n",
      "100%|█████████████████████████████████████████████████████████████████████████████████| 5039/5039 [00:00<00:00, 5616.90it/s]\n",
      "100%|█████████████████████████████████████████████████████████████████████████████████| 5039/5039 [00:00<00:00, 5613.44it/s]\n",
      "100%|█████████████████████████████████████████████████████████████████████████████████████| 27/27 [00:00<00:00, 3567.82it/s]\n",
      "100%|█████████████████████████████████████████████████████████████████████████████████| 5039/5039 [00:00<00:00, 5610.11it/s]\n",
      "100%|█████████████████████████████████████████████████████████████████████████████████| 5039/5039 [00:00<00:00, 5576.46it/s]\n",
      "100%|█████████████████████████████████████████████████████████████████████████████████| 5039/5039 [00:00<00:00, 5196.16it/s]\n",
      "100%|█████████████████████████████████████████████████████████████████████████████████| 5039/5039 [00:00<00:00, 5629.43it/s]\n",
      "100%|█████████████████████████████████████████████████████████████████████████████████| 5039/5039 [00:00<00:00, 5656.71it/s]\n",
      "100%|█████████████████████████████████████████████████████████████████████████████████| 5039/5039 [00:00<00:00, 5656.53it/s]\n",
      "100%|█████████████████████████████████████████████████████████████████████████████████| 5039/5039 [00:00<00:00, 5649.69it/s]\n",
      "100%|█████████████████████████████████████████████████████████████████████████████████| 5039/5039 [00:00<00:00, 5683.10it/s]\n",
      "100%|█████████████████████████████████████████████████████████████████████████████████| 5039/5039 [00:00<00:00, 5601.93it/s]\n",
      "100%|█████████████████████████████████████████████████████████████████████████████████| 5039/5039 [00:01<00:00, 4884.01it/s]\n",
      "100%|█████████████████████████████████████████████████████████████████████████████████| 5039/5039 [00:00<00:00, 5659.00it/s]\n"
     ]
    }
   ],
   "source": [
    "processed = multiprocessing(df, loop, cores = 30)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "863a8663",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "151197"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(processed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "ef75b660",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('prepare-MusicBench-Instructions.json', 'w') as fopen:\n",
    "    json.dump(processed, fopen)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "f3df919a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'text': '<|im_start|>system\\nYou are a helpful assistant.<|im_end|>\\n<|im_start|>user\\nAudio 1: <|audio_bos|><|AUDIO|><|audio_eos|>\\ncan u describe the audio<|im_end|>\\n<|im_start|>assistant\\nThis mellow instrumental track showcases a dominant electric guitar that opens with a descending riff, followed by arpeggiated chords, hammer-ons, and a slide. The percussion section keeps it simple with rim shots and a common time count, while the bass adds a single note on the first beat of every bar. Minimalist piano chords round out the song while leaving space for the guitar to shine. There are no vocals, making it perfect for a coffee shop or some chill background music. The key is in E major, with a chord progression that centers around that key and a straightforward 4/4 time signature.<|im_end|>\\n',\n",
       " 'audio': 'datashare/data_aug2/-0SdAVK79lg_1.wav'}"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "processed[0]"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
